import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import confusion_matrix, classification_report, roc_auc_score, mean_absolute_error, \
    mean_squared_error, accuracy_score
import dataPreproces
import json
from xgboostTrain import printJuZhen, saveJson


def mx_line(x_train, y_train):
    # 线性回归算法
    mx = LinearRegression()
    mx.fit(x_train, y_train)
    return mx


def mx_log(x_train, y_train):
    # 逻辑回归算法
    mx = LogisticRegression(penalty=12)
    mx.fit(x_train, y_train)
    return mx


def mx_bayes(x_train, y_train):
    # 多项式朴素贝叶斯算法
    mx = MultinomialNB(alpha=0.01)
    mx.fit(x_train, y_train)
    return mx


def mx_knn(x_train, y_train):
    # KNN邻近算法
    mx = KNeighborsClassifier()
    mx.fit(x_train, y_train)
    return mx


def mx_forest(x_train, y_train):
    # 随机森林算法
    mx = RandomForestClassifier(n_estimators=20)
    mx.fit(x_train, y_train)
    return mx


if __name__ == '__main__':
    # 获取数据并进行预处理
    data = dataPreproces.data_pretreat()
    # 卡方检测前13个特征
    feature_set = ['ResponseTimeTimeVariance', 'PacketLengthMedian', 'ResponseTimeTimeStandardDeviation',
                   'PacketLengthCoefficientofVariation', 'ResponseTimeTimeSkewFromMode', 'PacketLengthMode',
                   'ResponseTimeTimeSkewFromMedian', 'PacketTimeVariance', 'PacketTimeMedian', 'PacketTimeMode',
                   'PacketTimeStandardDeviation', 'PacketTimeMean', 'Duration']
    X = data[feature_set]
    Y = data['DoH']
    x_train, x_test, y_train, y_test = train_test_split(X, Y)
    # 生成模型mx
    mx = mx_forest(x_train, y_train)
    # 用模型mx预测x_test里的数据并返回结果集y_pred
    y_pred = mx.predict(x_test.values)
    # print(y_pred)
    print('---------------')
    # print(y_test)
    print('---------------')
    result = {'模型名称': '随机森林', '结果': {
        '准确率': accuracy_score(y_test, y_pred),
        '平均绝对误差': mean_absolute_error(y_test, y_pred),
        '均方误差': mean_squared_error(y_test, y_pred),
    }}
    saveJson(result, '../RandomForestResult.json')
    '''
        模型评估
        参考网站:https://blog.csdn.net/sinat_16388393/article/details/91427631
    '''
    target_names = ['True', 'False']
    labels = [0, 1]
    y_pred = [int(item > 0) for item in y_pred]
    print(classification_report(y_test, y_pred, labels=labels, target_names=target_names))
    # 混淆矩阵
    confusion_mat = confusion_matrix(y_test, y_pred)
    print(confusion_mat)
    printJuZhen(confusion_mat)
